"""阿里百炼平台QwQ模型提供商 - 支持深度思考模型"""

import os
from typing import List, Dict, Any, Optional, Generator
from openai import OpenAI
from dotenv import load_dotenv


class BailianQwQProvider:
    """阿里百炼平台QwQ模型提供商 - 支持深度思考模型"""
    
    def __init__(self, api_key: str = None):
        """初始化QwQ提供商"""
        load_dotenv()
        self.api_key = api_key or os.getenv("DASHSCOPE_API_KEY")
        self.client = OpenAI(
            api_key=self.api_key,
            base_url="https://dashscope.aliyuncs.com/compatible-mode/v1",
        )
    
    def stream_thinking_chat(self, messages: List[Dict[str, str]], model: str = "qwq-plus", **kwargs) -> Generator[Dict[str, Any], None, None]:
        """
        流式深度思考聊天接口
        
        支持模型返回两个阶段的内容：
        1. reasoning: 深度思考过程（逐步输出）
        2. answer: 最终回答（思考完成后开始输出）

        Args:
            messages: 聊天消息列表，格式为 [{"role": "user", "content": "..." }]
            model: 模型名称，默认为 "qwq-plus"
            **kwargs: 其他传递给 API 的参数（如 temperature, max_tokens 等）

        Yields:
            dict: 包含以下类型的事件：
                - {"type": "reasoning", "content": "..."}         # 思考中
                - {"type": "answer_start", "content": "..."}      # 回答开始
                - {"type": "answer", "content": "..."}            # 回答继续
                - {"type": "usage", "content": Usage Object}      # 用量信息（最后）
        """
        completion = self.client.chat.completions.create(
            model=model,
            messages=messages,
            extra_body={"enable_thinking": True},
            stream=True,
            **kwargs
        )
        
        reasoning_content = ""  # 完整思考过程
        answer_content = ""  # 完整回复
        is_answering = False  # 是否进入回复阶段
        
        for chunk in completion:
            if not chunk.choices:
                yield {
                    "type": "usage",
                    "content": chunk.usage
                }
                continue

            delta = chunk.choices[0].delta
            result = {"type": "chunk", "content": ""}

            # 只收集思考内容
            if hasattr(delta, "reasoning_content") and delta.reasoning_content is not None:
                if not is_answering:
                    result["type"] = "reasoning"
                reasoning_content += delta.reasoning_content
                result["content"] = delta.reasoning_content

            # 收到content，开始进行回复
            if hasattr(delta, "content") and delta.content:
                if not is_answering:
                    result["type"] = "answer_start"
                    is_answering = True
                else:
                    result["type"] = "answer"
                answer_content += delta.content
                result["content"] = delta.content
            
            yield result
    
    def get_thinking_models(self) -> List[str]:
        """获取支持深度思考的模型列表"""
        return [
            "qwq-plus",
            "qwen3-30b-a3b-thinking-2507",
            "qwen3-235b-a22b-thinking-2507"
        ]


def test_qwq_thinking():
    """测试QwQ深度思考功能"""
    try:
        provider = BailianQwQProvider()
        
        messages = [{"role": "user", "content": "你是谁"}]
        
        print("\n" + "=" * 20 + "思考过程" + "=" * 20 + "\n")
        
        for chunk in provider.stream_thinking_chat(messages):
            if chunk["type"] == "reasoning":
                print(chunk["content"], end="", flush=True)
            elif chunk["type"] == "answer_start":
                print("\n" + "=" * 20 + "完整回复" + "=" * 20 + "\n")
                print(chunk["content"], end="", flush=True)
            elif chunk["type"] == "answer":
                print(chunk["content"], end="", flush=True)
            elif chunk["type"] == "usage":
                print(f"\n\n使用情况: {chunk['content']}")
        
        print("\n\nQwQ深度思考测试完成！")
        
    except Exception as e:
        print(f"QwQ深度思考测试失败: {e}")


if __name__ == "__main__":
    test_qwq_thinking()